/*
 * Copyright 2010 Mark Southern, The Scripps Research Institute
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.googlecode.exceltablemodel;

import java.awt.image.RenderedImage;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.StringWriter;
import java.io.Writer;
import java.util.Calendar;
import java.util.Date;
import java.util.Iterator;

import javax.imageio.ImageIO;
import javax.swing.event.TableModelEvent;
import javax.swing.table.AbstractTableModel;

import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.apache.poi.openxml4j.exceptions.InvalidFormatException;
import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.DataFormatter;
import org.apache.poi.ss.usermodel.DateUtil;
import org.apache.poi.ss.usermodel.Drawing;
import org.apache.poi.ss.usermodel.ErrorConstants;
import org.apache.poi.ss.usermodel.FormulaEvaluator;
import org.apache.poi.ss.usermodel.RichTextString;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.ss.usermodel.WorkbookFactory;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;

/*
 * @author Mark Southern (southermark at gmail dot com)
 */
public class ExcelTableModel extends AbstractTableModel {

	public enum ValueType {
		ACTUAL(Object.class), CELL(Cell.class), STRING(String.class);
		private Class<?> clazz;

		private ValueType(Class<?> clazz) {
			this.clazz = clazz;
		}

		public Class<?> getTypeClass() {
			return clazz;
		}
	};
	
	public static ExcelTableModel newXLSInstance(boolean firstRowAsColumnHeadings) {
		return newInstance(new HSSFWorkbook(), firstRowAsColumnHeadings);
	}
	
	public static ExcelTableModel newXLSXInstance(boolean firstRowAsColumnHeadings) {
		return newInstance(new XSSFWorkbook(), firstRowAsColumnHeadings);
	}
	
	protected static ExcelTableModel newInstance(Workbook workbook, boolean firstRowAsColumnHeadings) {
		ExcelTableModel model = new ExcelTableModel();
		workbook.createSheet();
		model.setWorkbook(workbook);
		model.setUseFirstRowAsColumnHeadings(firstRowAsColumnHeadings);
		return model;
	}
	
	public static ExcelTableModel load(File file, boolean firstRowAsColumnHeadings) throws IOException, InvalidFormatException {
		return load(new FileInputStream(file), firstRowAsColumnHeadings);
	}	

	public static ExcelTableModel load(InputStream inputStream, boolean firstRowAsColumnHeadings) throws IOException, InvalidFormatException {
		Workbook workbook = WorkbookFactory.create(inputStream);
		ExcelTableModel model = new ExcelTableModel();
		model.setWorkbook(workbook);
		model.setUseFirstRowAsColumnHeadings(firstRowAsColumnHeadings);
		return model;
	}
	
	private Drawing drawing = null;
	private Integer maxColumns = 0;
	private Sheet sheet = null;
	private boolean useFirstRowAsColumnHeadings = false;
	private ValueType valueType = ValueType.CELL;

	private Workbook workbook;
	public ExcelTableModel() {

	}

	public Cell getCellAt(int row, int col) {
		if (isUseFirstRowAsColumnHeadings())
			row++;
		Row r = getSheet().getRow(row);
		if( r == null)
			r = getSheet().createRow(row);
		Cell cell = r.getCell((short) col);
		if( cell == null )
			cell = r.createCell(col);
		return cell;
	}

	public Object getCellValue(Cell cell) {
		return cell == null ? null : getCellValue(cell, cell.getCellType());
	}

	protected Object getCellValue(Cell cell, int type) {
		Object obj = null;
		if (type == Cell.CELL_TYPE_BOOLEAN)
			obj = cell.getBooleanCellValue();
		else if (type == Cell.CELL_TYPE_ERROR)
			obj = ErrorConstants.getText(cell.getErrorCellValue());
		else if (type == Cell.CELL_TYPE_NUMERIC) {
			if (DateUtil.isCellDateFormatted(cell))
				obj = cell.getDateCellValue();
			else
				obj = cell.getNumericCellValue();
		} else if (type == Cell.CELL_TYPE_STRING)
			obj = cell.getStringCellValue();
		else if (type == Cell.CELL_TYPE_BLANK)
			obj = null;
		else if (type == Cell.CELL_TYPE_FORMULA)
			obj = getCellValue(cell, cell.getCachedFormulaResultType());
		return obj;
	}

	public String getCellValueAsString(int row, int col) {
		DataFormatter formatter = new DataFormatter();
		FormulaEvaluator eval = getWorkbook().getCreationHelper().createFormulaEvaluator();
		String value;
		Cell cell = getCellAt(row, col);
		if( null == cell )
			value = null;
		else if (cell.getCellType() == Cell.CELL_TYPE_FORMULA)
			if (cell.getCachedFormulaResultType() == Cell.CELL_TYPE_ERROR)
				value = ErrorConstants.getText(cell.getErrorCellValue());
			else
				value = formatter.formatCellValue(cell, eval);
		else
			value = formatter.formatCellValue(cell);
		return value;
	}

	public Object getCellValueAt(int row, int col) {
		Cell cell = getCellAt(row, col);
		return getCellValue(cell);
	}

	@Override
	public Class<?> getColumnClass(int c) {
		return getValueType().getTypeClass();
	}

	@Override
	public int getColumnCount() {
		return this.maxColumns + 1;
	}

	@Override
	public String getColumnName(int column) {
		if (isUseFirstRowAsColumnHeadings()) {
			// row -1 b/c in this case we increment the row in getCellAt()
			String name = getCellValueAsString(-1, column);
			return name;
		} else
			return super.getColumnName(column);
	}
	
	public void setColumnName(int column, String name) {
		if (isUseFirstRowAsColumnHeadings()) {
			// row -1 b/c in this case we increment the row in getCellAt()
			setValueAt(name, -1, column);
		}
	}

	public Drawing getDrawing() {
		return drawing;
	}

	@Override
	public int getRowCount() {
		int rows = sheet.getLastRowNum() + 1;
		if (isUseFirstRowAsColumnHeadings())
			rows--;
		return rows;
	}

	public Sheet getSheet() {
		return sheet;
	}

	@Override
	public Object getValueAt(int row, int col) {
		Object obj = null;
		if (ValueType.CELL.equals(getValueType()))
			obj = getCellAt(row, col);
		else if (ValueType.ACTUAL.equals(getValueType()))
			obj = getCellValueAt(row, col);
		else
			obj = getCellValueAsString(row, col);
		return obj;
	}

	public ValueType getValueType() {
		return valueType;
	}

	public Workbook getWorkbook() {
		return workbook;
	}

	public boolean isCellEditable(int rowIndex, int columnIndex) {
		return true;
	}

	public boolean isUseFirstRowAsColumnHeadings() {
		return useFirstRowAsColumnHeadings;
	}

	public void setSheet(int sheetNumber) {
		Sheet sheet = getWorkbook().getSheetAt(sheetNumber);
		setSheet(sheet);
	}

	public void setSheet(Sheet sheet) {
		this.sheet = sheet;
		this.drawing = sheet.createDrawingPatriarch();
		Iterator<Row> iter = sheet.rowIterator();
		while (iter.hasNext()) {
			Row row = iter.next();
			int lastCol = row.getLastCellNum() - 1;
			if (maxColumns < lastCol)
				this.maxColumns = lastCol;
		}
		fireTableStructureChanged();
	}

	public void setSheet(String sheetName){
		Sheet sheet = getWorkbook().getSheet(sheetName);
		if( sheet == null )
			throw new IllegalArgumentException(String.format("Workbook does not contain a sheet named '%s'", sheetName));
		setSheet(sheet);
	}
	
	public void setUseFirstRowAsColumnHeadings(boolean useFirstRowAsColumnHeadings) {
		this.useFirstRowAsColumnHeadings = useFirstRowAsColumnHeadings;
		fireTableStructureChanged();
	}

	public void setValueAt(Object obj, int rowIndex, int columnIndex) {
		Cell cell = getCellAt(rowIndex, columnIndex);

		if (obj != null) {
			Class<?> clazz = obj.getClass();
			if (String.class.isAssignableFrom(clazz))
				cell.setCellValue(((String) obj));
			else if (Number.class.isAssignableFrom(clazz))
				cell.setCellValue(((Number) obj).doubleValue());
			else if (Boolean.class.isAssignableFrom(clazz))
				cell.setCellValue(((Boolean) obj).booleanValue());
			else if (Calendar.class.isAssignableFrom(clazz))
				cell.setCellValue(((Calendar) obj));
			else if (Date.class.isAssignableFrom(clazz))
				cell.setCellValue(((Date) obj));
			else if (RichTextString.class.isAssignableFrom(clazz))
				cell.setCellValue(((RichTextString) obj));
			else if (RenderedImage.class.isAssignableFrom(clazz)) {
				try {
					RenderedImage img = (RenderedImage)obj;
					ByteArrayOutputStream baos = new ByteArrayOutputStream();
					ImageIO.write(img, "png", baos);
					new ExcelImageUtils().insertCellImage(cell, getDrawing(), baos.toByteArray(), Workbook.PICTURE_TYPE_PNG, img.getWidth(), img.getHeight());
				}
				catch(IOException ex) {
					throw new RuntimeException(ex);
				}
			}
			else
				cell.setCellValue(obj.toString());
		}

		if (columnIndex > this.maxColumns) {
			this.maxColumns = columnIndex;
			fireTableStructureChanged();
		} else
			fireTableChanged(new TableModelEvent(this, rowIndex, rowIndex, columnIndex));
	}

	public void setValueType(ValueType cellType) {
		this.valueType = cellType;
	}

	public void setWorkbook(Workbook workbook) {
		this.workbook = workbook;
		setSheet(workbook.getActiveSheetIndex());
	}
	
	public Object[] getRow(int row) {
		Object[] objs = new Object[getColumnCount()];
		for(int col = 0; col < getColumnCount(); col++) {
			objs[col] = getValueAt(row, col);
		}
		return objs;
	}

	public String toString() {
		String s = "";
		try {
			StringWriter sw = new StringWriter();
			writeDelimited(sw, "\t", true);
			s = sw.toString();
		} catch (IOException ex) {
			ex.printStackTrace();
		}
		return s;
	}

	public void writeDelimited(Writer writer, String delimiter, boolean columnHeadings) throws IOException {
		if (columnHeadings) {
			for (int col = 0; col < getColumnCount(); col++) {
				String name = getColumnName(col);
				writer.write(name == null ? "" : name);
				if (col < getColumnCount() - 1)
					writer.write(delimiter);
			}
			writer.write(System.getProperty("line.separator"));
		}
		for (int row = 0; row < getRowCount(); row++) {
			for (int col = 0; col < getColumnCount(); col++) {
				String value = getCellValueAsString(row, col);
				writer.write(value == null ? "" : value);
				if (col < getColumnCount() - 1)
					writer.write(delimiter);
			}
			writer.write(System.getProperty("line.separator"));
		}
		writer.flush();
	}
}